#include "inputPprIdid.h"
#include "IdidEnterModels.h"
#include "ExpandTimeSteps.h"

//////////////////////////////////////////////////////////////////////////
//Author	:	Ross Conroy ross.conroy@tees.ac.uk
//Date		:	29/07/2014
//
//This class is responsible for taking each PPR tree and adding it to the
//model nodes of an I-DID using the following pseudo-code
//
//For Each model node
//	set number of states = number of models
//End For
//
//For Each TimeStep
//	generate combinations of parent values
//		For each model
//			For each combination in combinations
//				For each node in combination
//					If exists in model 
//						If is end node
//							action = model action
//						Else
//							move to next node
//						End if
//					Else
//						action = all actions with even probability
//					EndIf
//				End for
//				Apply action(s) to model nodes	
//			End For
//		End For
//End For
//
//Update 10/08/2014
//Re write to better support generic I-DIDs and change of algorithm
//
//For Each actionNode IN agent j actions
//	For Each parent state In parent.getStates
//		Generate PPR branch using parent state name comma separated string
//		First item in string indicates the model
//
//		For Each PPR in PPRs
//			If Branch exists in PPR
//				Record action to list
//			Else
//				Do Nothing
//			End If
//		End For
//
//		Count ocurances of each action and record in CPT for current action node
//
//	End For
//EndFor
//////////////////////////////////////////////////////////////////////////

InputPprIdid::InputPprIdid(void)
{
	time_t timeVal=0;
	time(&timeVal);
	srand(timeVal);
}

//////////////////////////////////////////////////////////////////////////
//Executes the above pseudo-code
//////////////////////////////////////////////////////////////////////////
Domain * InputPprIdid::InputPprIntoIdidModels(Domain * inputID, hash_map<string, PPR> pprs, string jActionPrefix)
{
	return inputID;
}

Domain * InputPprIdid::EnterPprModelsIntoIdid(Domain * inputIDID, hash_map<string, PPR> pprs, string modelPrefix, string jActionPrefix, int NumModels, bool randomUnknown)
{
	IdidEnterModels ididentermodels;

	//loop through and get model name for each from root node
	vector<string> modelNames;
	for(hash_map<string, PPR>::iterator it = pprs.begin(); it != pprs.end(); it++)
	{
		PPR tempPPR = it->second;
		string tempStr = tempPPR.GetRootNode()->GetState();
		string tempStr2 = tempStr.substr (0,tempStr.length()-1);
		modelNames.push_back(tempStr2);
	}

	inputIDID = ididentermodels.ExpandModelNodesSteps(inputIDID, modelPrefix, NumModels, modelNames);

	ExpandTimeSteps expander;
	int numSteps =  expander.FindLastTimeStep(inputIDID);

	for(int i = 1; i <= numSteps; i++)
	{
		stringstream ss;
		ss << jActionPrefix;
		ss << i;
		DiscreteNode * modelActionNode = (DiscreteNode *)inputIDID->getNodeByName(ss.str());
		DiscreteNode * parent = (DiscreteNode *)modelActionNode->getParents().front();

		NumberList nl = modelActionNode->getTable()->getData();

		int tableIndex = 0;

		for(int s = 0; s < parent->getNumberOfStates(); s++)
		{
			vector<int> ocurances (modelActionNode->getNumberOfStates());
			int rState = rand() % modelActionNode->getNumberOfStates();
			//loop through each PPR to find if any match the search string
			for (hash_map<string, PPR>::iterator itp = pprs.begin(); itp != pprs.end(); itp++)
			{
				PPR ppr = itp->second;
				OcurancesOfAction ocurancesAction = ppr.GetOcurancesOfActionFromSearchString(parent->getStateLabel(s));

				//if found something then add it to the vector of ocurances
				if(ocurancesAction.ocurances > 0)
				{
					int stateIndex = modelActionNode->getStateIndex(ocurancesAction.action);
					ocurances[stateIndex] += ocurancesAction.ocurances;
				}
				else if (ocurancesAction.ocurances == -1)
				{ 
					//Random					
					ocurances[rState] += 100;
				}
			}

			//check that all of the vector is != 0
			bool valid = false;
			{
				for(int v = 0; (v < ocurances.size()) && !valid; v++)
				{
					if(ocurances[v] > 0)
					{
						valid = true;
					}
				}
			}
			if(!valid)
			{
				if(false)
				{
					int rState = rand() % ocurances.size();
					ocurances[rState] = 1;
				}
				else
				{
					for(int v = 0; (v < ocurances.size()) && !valid; v++)
					{
						ocurances[v] = 1;
					}
				}				
			}

			for(int v = 0; v < ocurances.size(); v++)
			{
				nl[tableIndex] = ocurances[v];
				tableIndex ++;
			}
		}

		modelActionNode->getTable()->setData(nl);
	}

	return inputIDID;
}


//////////////////////////////////////////////////////////////////////////
//generates and loops through all potential parent combinations for each
//node to determine if the model can be applied to each combination
//////////////////////////////////////////////////////////////////////////
void InputPprIdid::ApplyPprToNode(DiscreteNode * node, hash_map<string, PPR> pprsIn)
{
	hash_map<int, PPR> pprs;
	int i = 0;
	for(hash_map<string, PPR>::iterator it = pprsIn.begin(); it!= pprsIn.end(); it++)
	{
		PPR tempPPR = it->second;
		pprs.insert(make_pair(i, tempPPR));
		i++;
	}

	list<int> sizes;	
	NodeList parents = node->getParents();
	for(NodeList::const_iterator itp = parents.begin(); itp != parents.end(); itp++)
	{
		DiscreteNode * node = (DiscreteNode*)*itp;
		sizes.push_back(node->getNumberOfStates());
	}
	CombinationGenerator cg;
	list<list<int>> combinations;
	combinations = cg.GenerateCombinations(sizes);

	NumberList table = node->getTable()->getData();
	int width = table.size() / node->getNumberOfStates();
	for(NumberList::iterator itnl = table.begin(); itnl != table.end(); itnl++)
	{
		*itnl = 0;
	}

	for (list<list<int>>::iterator it = combinations.begin(); it != combinations.end(); it++)
	{
		list<int> combination = *it;
		PPR model = pprs.find(combination.front())->second;
		list<string> actions = GetActionsForCombinationAndPpr(combination, model, node, true);

		//now set actions in cpt from combinations and actions list
		int i = 0;
		list<int> divs;
		int divsum = 1;
		int pos = 0;

		for(list<int>::iterator itc = combination.begin(); itc !=combination.end(); ++itc)
		{
			DiscreteNode * parent = (DiscreteNode *)parents[i];
			int parentState = *itc;
			divs.push_back(parent->getNumberOfStates());
			divsum = divsum * divs.back();
			pos += ((width / divsum) * (parentState));
			i++;
		}

		for (int j = 0; j < node->getNumberOfStates(); j++)
		{
			for(list<string>::iterator ita = actions.begin(); ita != actions.end(); ita++)
			{
				string stateName = *ita;
				if(node->getStateLabel(j) == stateName)
				{
					int index = ((int)node->getNumberOfStates() * pos) + (int)j;
					table[index] = 1;
				}
			}							
		}

	}

	node->getTable()->setData(table);
}

//////////////////////////////////////////////////////////////////////////
//for a given combination and model will return what the expected action
//should be
//////////////////////////////////////////////////////////////////////////
list<string> InputPprIdid::GetActionsForCombinationAndPpr(list<int> combination, PPR model, DiscreteNode * node, bool randomUnknown)
{
	list<string> actions;
	list<string> searchBranch;

	NodeList parents = node->getParents();
	NodeList::const_iterator itp = parents.begin();
	itp++;

	for(list<int>::iterator itc = combination.begin(); itc != combination.end(); itc++)
	{
		if(itc != combination.begin())
		{
			DiscreteNode * parent = (DiscreteNode*)*itp;
			int parentStateIndex = *itc;

			string parentState = parent->getStateLabel(parentStateIndex);
			searchBranch.push_back(parentState);

			itp++;
		}
	}

	//Search for combination in tree
	ActionNode * currentNode = model.GetRootNode();
	bool match = true;

	for(list<string>::iterator it = searchBranch.begin(); (it != searchBranch.end()) && (match == true); it++)
	{
		string actionName = *it;

		if(currentNode->GetNodeName() == actionName)
		{
			it++;
			string obs = *it;
			currentNode = currentNode->GetChildNodes().find(obs)->second;
			
			if(currentNode == NULL) //will be null when an empty branch is found
			{
				match = false;
			}
		}
		else
		{
			match = false;
		}
	}

	//if a match was found then current action will indicate the expected action
	if(match)
	{
		actions.push_back(currentNode->GetNodeName());
	}
	else if(randomUnknown)
	{
		//If unknown and allowed to randomize unknown values
		int rState = rand() % node->getNumberOfStates();
		actions.push_back(node->getStateLabel(rState));		
		cout << "Random " << node->getStateLabel(rState) << endl;
	}
	else
	{
		for(int i = 0; i < node->getNumberOfStates(); i++)
		{
			actions.push_back(node->getStateLabel(i));
		}
	}


	return actions;
}

//////////////////////////////////////////////////////////////////////////
//Loops through all of the model nodes in the domain and sets the number
//of states for each to the specified number of models
//////////////////////////////////////////////////////////////////////////
Domain * InputPprIdid::SetModelNodesStates(Domain * inputID, string modelPrefix, int numModels)
{
	//Must be a positive integer
	if(numModels < 1)
	{
		return inputID;
	}

	NodeList nodes = inputID->getNodes();
	NodeList modelNodes;

	for (NodeList::const_iterator it = nodes.begin(); it != nodes.end(); it++)
	{
		DiscreteNode * node = (DiscreteNode*)*it;
		string nodeName = node->getName();

		if (nodeName.find(modelPrefix) != std::string::npos) 
		{
			modelNodes.push_back(node);
		}
	}

	//set states
	for (NodeList::const_iterator it = modelNodes.begin(); it != modelNodes.end(); it++)
	{
		DiscreteNode * node = (DiscreteNode*)*it;
		node->setNumberOfStates(numModels);
		for(int i = 0; i < numModels; i++)
		{
			stringstream ss;
			ss << "Model" << i;
			node->setStateLabel(i, ss.str());
		}
	}

	//set CPTs to diagonal 1's
	for (NodeList::const_iterator it = modelNodes.begin(); it != modelNodes.end(); it++)
	{
		DiscreteNode * node = (DiscreteNode*)*it;
		NumberList table = node->getTable()->getData();

		if(table.size() != node->getNumberOfStates())
		{
			for(int i = 0; i < table.size(); i++)
			{
				table[i] = 0;
			}

			for(int i = 0; i < table.size(); i += (node->getNumberOfStates() + 1))
			{
				table[i] = 1;
			}			
		}	
		else
		{
			for(int i = 0; i < table.size(); i++)
			{
				table[i] = 1;
			}
		}

		node->getTable()->setData(table);
	}

	return inputID;
}